package com.amumu.drama.common.commpont;


import cn.hutool.core.lang.generator.SnowflakeGenerator;
import com.amumu.drama.common.base.annonation.OnContextInited;
import com.amumu.drama.common.config.redis.AppRateLimitType;
import com.amumu.drama.common.config.redis.JedisRateLimiterConfig;
import com.amumu.drama.common.config.redis.JedisTemplate;
import com.amumu.drama.common.config.redis.RateLimitConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.lang.management.ManagementFactory;

/**
 * @author sunzhuang
 * @date 2024/9/12 15:41
 * @slogan 天下风云出我辈, 一入代码岁月催
 * @desc todo
 */
@Component
public class RateLimiterCache {

    public static RateLimiterCache INSTANT;

    public static String MACHINE_NAME;

    @Autowired
    private JedisTemplate jedisTemplate;

    @OnContextInited
    public void init() {
        INSTANT = SpringContext.getBean(RateLimiterCache.class);
        MACHINE_NAME = System.getProperty("app.id") + ":" + getHostName();
    }

    private static String getHostName() {
        String hostName = new SnowflakeGenerator().next() + "";
        String pidAtHostName = ManagementFactory.getRuntimeMXBean().getName();
        int idx = pidAtHostName.indexOf('@');
        if (idx > 0) {
            hostName = pidAtHostName.substring(idx + 1);
        }
        return hostName;
    }

    /** 根据Key限流 */
    public boolean allowByKey(String key, AppRateLimitType limitType) {
        JedisRateLimiterConfig cfg = RateLimitConfig.RATE_LIMITER_CONFIG_MAP.get(limitType);
        if (cfg == null || cfg.getLimitCount() < 1) {
            return true;
        }
        try {
            return jedisTemplate.rateLimit(cfg).tryAcquire(key);
        } catch (Exception e) {
            return RateLimitConfig.PASS_IF_ERROR == 1;
        }
    }
}
